from gurobipy import *
import numpy as np
from traffic_engineer import *
from robust_multi_cluster_traffic_engineer import *
from traffic_engineering_class import *

class HistoricalAveTrafficEngineer(TrafficEngineer):
	## Takes in two new variables:
	## update_frequency - how much time in-between every traffic engineering events
	def __init__(self, aurora_network, all_interblock_paths, update_frequency, training_length):
		TrafficEngineer.__init__(self, aurora_network, all_interblock_paths)
		self.update_frequency = update_frequency
		self.training_length = training_length
		return

	def get_filename_param(self):
		return "avete_r{}t{}".format(self.update_frequency, self.training_length)

	# returns a string description of this class
	def get_string(self):
		return "Ave TE, w/ update_frequency: {} and training_length: {}".format(self.update_frequency, self.training_length)

	def traffic_engineering_class(self):
		return TrafficEngineeringClass.AVE_TE

	#def _compute_mlu(self, traffic_matrices, capacity_matrix, all_paths):
	## Currently does not involve another step that reduces hop count
	def _compute_path_weights_minimize_worstcase_mlu(self, logical_topology_adj_matrix, traffic_matrix):
		nblocks = self.aurora_network.get_num_blocks()
		model = Model("Routing minimize MLU for all traffic matrices")
		model.setParam( 'OutputFlag', False )
		mlu = model.addVar(lb=0., ub=GRB.INFINITY, obj=0, vtype=GRB.CONTINUOUS, name="mlu")
		routing_weight_vars = {}
		link_capacity_constraints = [None] * nblocks
		for i in range(nblocks):
			link_capacity_constraints[i] = [None] * nblocks
			for j in range(nblocks):
				if i != j:
					link_capacity_constraints[i][j] = LinExpr()
					#routing_weight_vars[path] = model.addVar(lb=0., ub=1., obj=1., vtype=GRB.CONTINUOUS, name="w_{},{}".format(i,j))
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					weight_constraint = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weight_vars[path] = model.addVar(lb=0., ub=1., obj=1., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						weight_constraint += routing_weight_vars[path]
					model.addConstr(lhs=weight_constraint, sense=GRB.EQUAL, rhs=1.)


		# stage 1: setup all the optimization variables that are the routing weights
		# stage 1.5 : also adds the traffic flow satisfiability constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					for path in self.all_interblock_paths[i][j]:
						src = path[0]
						for path_hop in range(1, len(path), 1):
							dst = path[path_hop]
							link_capacity_constraints[src][dst] += (routing_weight_vars[path] * traffic_matrix[i][j])
							src = dst
		## stage 2 : add the link utilization constraints to the model
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					lu_constraint = link_capacity_constraints[i][j]
					capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
					model.addConstr(lhs=lu_constraint, sense=GRB.LESS_EQUAL, rhs=mlu * capacity * logical_topology_adj_matrix[i][j])
		# stage 3: set the objective
		model.setObjective(mlu, GRB.MINIMIZE)
		try: 
			model.optimize()
			routing_weights = {}
			for path in routing_weight_vars.keys():
				routing_weights[path] = routing_weight_vars[path].x
			return routing_weights
		except GurobiError as e:
			print ("Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError :
			print ("Encountered an attribute error ")
			return None

	def compute_path_weights(self, logical_topology_adj_matrix, traffic_matrices):
		nblocks = self.aurora_network.get_num_blocks()
		expected_traffic_matrix = np.zeros((nblocks, nblocks))
		num_tms = len(traffic_matrices)
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					for tm_index in range(num_tms):
						expected_traffic_matrix[i][j] += (traffic_matrices[tm_index][i][j] / num_tms)
		routing_weights = self._compute_path_weights_minimize_worstcase_mlu(logical_topology_adj_matrix, expected_traffic_matrix)
		return routing_weights